from datetime import datetime, timedelta
from sqlalchemy import func
from app.models.transactions import Transaction
from app.models.budgets import Budget, CategoryBudget

class FinanceDataCollector:
    def __init__(self, db_session, user_id):
        self.db_session = db_session
        self.user_id = user_id
        # 默认取最近6个月的数据
        self.end_date = datetime.now()
        self.start_date = self.end_date - timedelta(days=180)
        
    def collect_basic_data(self):
        """收集基础财务数据"""
        # 获取收入和支出数据
        income_data = self._get_income_data()
        expense_data = self._get_expense_data()
        budget_data = self._get_budget_data()
        
        return {
            "time_range": {
                "start_date": self.start_date.strftime("%Y-%m-%d"),
                "end_date": self.end_date.strftime("%Y-%m-%d")
            },
            "income": income_data,
            "expense": expense_data,
            "budget": budget_data,
            "category_distribution": self._get_category_distribution()
        }
    
    def _get_income_data(self):
        """获取收入数据"""
        monthly_income = self.db_session.query(
        func.date_format(Transaction.date, '%Y-%m-01').label('month'),  # 使用 DATE_FORMAT 替换 date_trunc
        func.sum(Transaction.amount).label('total')
        ).filter(
            Transaction.user_id == self.user_id,
            Transaction.type == 'income',
            Transaction.date >= self.start_date,
            Transaction.date <= self.end_date
        ).group_by('month').order_by('month').all()
            
        return {
            "total": sum(item.total for item in monthly_income) if monthly_income else 0,
            "monthly": [{"month": item.month, "amount": float(item.total)} 
                       for item in monthly_income]
        }
    
    def _get_expense_data(self):
        """获取支出数据"""
        monthly_expense = self.db_session.query(
        func.date_format(Transaction.date, '%Y-%m-01').label('month'),  # 使用 DATE_FORMAT 替换 date_trunc
        func.sum(Transaction.amount).label('total')
        ).filter(
            Transaction.user_id == self.user_id,
            Transaction.type == 'expense',
            Transaction.date >= self.start_date,
            Transaction.date <= self.end_date
        ).group_by('month').order_by('month').all()
        
        return {
            "total": sum(item.total for item in monthly_expense) if monthly_expense else 0,
            "monthly": [{"month": item.month, "amount": float(item.total)} 
                       for item in monthly_expense]
        }
    
    def _get_budget_data(self):
        """获取预算数据"""
        # 获取月度预算
        monthly_budgets = self.db_session.query(
            Budget.month,
            Budget.year,
            Budget.total_budget
        ).filter(
            Budget.user_id == self.user_id,
            Budget.type == 'month',
            Budget.year == self.end_date.year
        ).all()
        
        return {
            "monthly": [{"year": item.year, "month": item.month, "amount": float(item.total_budget)}
                       for item in monthly_budgets]
        }
    
    def _get_category_distribution(self):
        """获取支出分类分布"""
        category_expense = self.db_session.query(
            Transaction.category_id,
            func.sum(Transaction.amount).label('total'),
            func.avg(Transaction.amount).label('avg_amount')
        ).filter(
            Transaction.user_id == self.user_id,
            Transaction.type == 'expense',
            Transaction.date >= self.start_date,
            Transaction.date <= self.end_date
        ).group_by(Transaction.category_id).all()
        
        return [{"category_id": item.category_id, "total": float(item.total), "avg": float(item.avg_amount)}
                for item in category_expense]